import torch
import numpy as np
from sklearn.model_selection import train_test_split


def np_to_torch(npdata, dtype="float"):
    """
    transform numpy array to torch tensor
    :param dtype: 数据类型，默认为float32
    :param npdata: a numpy array
    :return: a torch tensor
    """
    data = npdata.astype(dtype)
    data = torch.from_numpy(data)
    # torch_data = torch.tensor(data, dtype=torch.float32)
    if dtype == "float":
        torch_data = data.clone().detach().float()
    else:
        torch_data = data.clone().detach().double()
    return torch_data


def load_np_data(train_data_path, test_data_path, valid_split=0.2, random_state=0, dtype="float"):
    """
    加载numpy数据来训练剩余时间预测模型
    :param train_data_path:
    :param test_data_path:
    :param valid_split:
    :param random_state:
    :param dtype:
    :return:
    """
    train_data = np.load(train_data_path)
    test_data = np.load(test_data_path)
    train_xy, valid_xy = train_test_split(train_data, test_size=valid_split, shuffle=True, random_state=random_state)
    train_xy_torch = np_to_torch(train_xy, dtype)
    valid_xy_torch = np_to_torch(valid_xy, dtype)
    test_xy_torch = np_to_torch(test_data, dtype)
    xtrain = train_xy_torch[:, :, 1:-1]
    ytrain = torch.reshape(train_xy_torch[:, -1, -1], (train_xy_torch.shape[0], 1))
    xvalid = valid_xy_torch[:, :, 1:-1]
    yvalid = torch.reshape(valid_xy_torch[:, -1, -1], (valid_xy_torch.shape[0], 1))
    xtest = test_xy_torch[:, :, 1:-1]
    ytest = torch.reshape(test_xy_torch[:, -1, -1], (test_xy_torch.shape[0], 1))
    return xtrain, ytrain, xvalid, yvalid, xtest, ytest


def load_eval_data(train_data_path, test_data_path, dtype="float64"):
    """
    加载数据来生成预测数据
    :param train_data_path:
    :param test_data_path:
    :param dtype:
    :return:
    """
    train_data = np.load(train_data_path)
    test_data = np.load(test_data_path)
    train_xy_torch = np_to_torch(train_data, dtype)
    test_xy_torch = np_to_torch(test_data, dtype)
    xtrain = train_xy_torch[:, :, :-1]
    ytrain = torch.reshape(train_xy_torch[:, -1, -1], (train_xy_torch.shape[0], 1))
    xtest = test_xy_torch[:, :, :-1]
    ytest = torch.reshape(test_xy_torch[:, -1, -1], (test_xy_torch.shape[0], 1))
    return xtrain, ytrain, xtest, ytest


def bpic2013_abnormal_data_clean():
    abnormal_cid = [740542029]
    test_np_data = np.load("data/drl_data/bpic2013/test_case_ids.npy")
    clean_cids = []
    for cid in test_np_data:
        if cid not in abnormal_cid:
            clean_cids.append(cid)
    np.save("data/drl_data/bpic2013/test_case_ids_v1.npy", np.array(clean_cids))
    print(len(clean_cids))


def bpic2018_abnormal_data_clean():
    abnormal_cid = [-1411132132, 1311396517]
    test_np_data = np.load("data/drl_data/bpic2018/test_case_ids.npy")
    clean_cids = []
    for cid in test_np_data:
        if cid not in abnormal_cid:
            clean_cids.append(cid)
    np.save("data/drl_data/bpic2018/test_case_ids_v1.npy", np.array(clean_cids))
    print(len(clean_cids))


if __name__ == '__main__':
    bpic2018_abnormal_data_clean()
